import tkinter as tk
from PIL import Image, ImageTk
import numpy as np
import socket
import threading

# UDP 配置
udp_ip = "127.0.0.1"  # 替换为接收端的 IP 地址
udp_port = 12345
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.bind((udp_ip, udp_port))

# 图像尺寸配置
width, height = 640, 480

class App:
    def __init__(self, root):
        self.root = root
        self.root.title("UDP 图像接收")
        self.label = tk.Label(root)
        self.label.pack()
        self.thread = threading.Thread(target=self.receive_data)
        self.thread.daemon = True
        self.thread.start()

    def receive_data(self):
        buffer = bytearray(width * height * 2)  # 创建一个足够大的缓冲区
        while True:
            # 接收数据
            data, _ = sock.recvfrom(1280)
            buffer.extend(data)
            
            # 如果缓冲区满了，处理图像
            if len(buffer) >= width * height * 2:
                # 转换为 numpy 数组
                img_rgb565 = np.frombuffer(buffer[:width * height * 2], dtype=np.uint16).reshape(height, width)
                
                # 清空缓冲区
                buffer = buffer[width * height * 2:]
                
                # 解码 RGB565 为 RGB888
                r = (img_rgb565 >> 8) & 0xF8
                g = (img_rgb565 >> 3) & 0xFC
                b = (img_rgb565 << 3) & 0xF8
                img_rgb888 = np.stack((r, g, b), axis=-1).astype(np.uint8)
                
                # 转换为 PIL 图像
                img = Image.fromarray(img_rgb888, 'RGB')
                
                # 更新图像到 GUI
                self.update_image(img)

    def update_image(self, img):
        img_tk = ImageTk.PhotoImage(image=img)
        self.label.config(image=img_tk)
        self.label.image = img_tk  # 保持引用，防止被垃圾回收

if __name__ == "__main__":
    root = tk.Tk()
    app = App(root)
    root.mainloop()
